#!/usr/bin/python3
#
# Copyright (C) 2024, 2025, Cyrus Li (e21920076@126.com). All rights reserved.
#

"""
Create ssh config block
"""

import sys
import yaml


USER_BLACKLIST_PM = ['anuser', 'tone']

DOMAIN_USERS_YAML = """
- user: root
  password: Passw0rd
- user: tone
  password: tone
- user: anuser
  password: anolisos
"""

DOMAIN_PMS_YAML = """
- iid: 0x1
  alias: pc
  ip: 192.168.100.13
  arch: x86_64
  user: root
  password: OrzLJ@69

- iid: 0x2
  alias: pa
  ip: 192.168.100.11
  arch: aarch64
  user: root
  password: OrzLJ@69
"""

DOMAIN_VMS_YAML = """
- iid: 1001
  alias: va
  ip: 191.168.122.161
  os: Anolis OS 23.1
  role: T-One master
  pmid: 0x1

- iid: 1003
  alias: vc
  ip: 191.168.122.163
  os: Fedora 41
  role: SUT
  pmid: 0x2
"""

DOMAIN_VM_IID_BASE = 1000


def load_yaml(text):
    obj = yaml.safe_load(text)
    return obj


def get_pm_byid(dom_pms, pmid):
    for pm in dom_pms:
        if pmid == pm.get('iid'):
            return pm
    return None


def output_one_block(host, hostname, user,
                     *,
                     password=None,
                     os=None,
                     role=None,
                     ip=None,
                     arch=None):
    print("#" * 99)
    print(f"Host\t\t{host}")
    print(f"HostName\t{hostname}")
    print(f"User\t\t{user}")
    print(f"#Password\t{password}") if password is not None else print(end='')
    print(f"#OS\t\t{os}")           if os       is not None else print(end='')
    print(f"#Role\t\t{role}")       if role     is not None else print(end='')
    print(f"#Arch\t\t{arch}")       if arch     is not None else print(end='')
    print(f"#PM\t\t{ip}")           if ip       is not None else print(end='')
    print(f"#SSH:\t\tssh {host.split()[0]}\t# Password: {password}\t# i.e. ssh {user}@{hostname}")
    print()


def get_host_aliases(hostname, alias, ip, user):
    out = []
    out.append(f"{alias}{user[0]}")
    out.append(f"{user}{hostname}")
    out.append(f"{user}-{ip}")
    return ' '.join(out)


def patch_dom_users(dom_users, user_blacklist, user, password):
    users = []
    for dom_user in dom_users:
        if user_blacklist is not None and dom_user.get('user') in user_blacklist:
            continue

        if dom_user.get('user') == user:
            users.append({'user': user, 'password': password})
        else:
            users.append(dom_user)
    return users


def dump_ssh_cfg_block_pm(dom_pms, dom_users):
    pms_num = len(dom_pms)
    pms_hostname_width = 3 if len(str(pms_num)) < 3 else len(str(pms_num))

    for pm in dom_pms:
        pm_iid = pm.get('iid')
        pm_alias = pm.get('alias')
        pm_ip = pm.get('ip')
        pm_arch = pm.get('arch')
        pm_hostname = f"PM{pm_iid:0{pms_hostname_width}}"
        for dom_user in patch_dom_users(dom_users, USER_BLACKLIST_PM, pm.get('user'), pm.get('password')):
            user = dom_user.get('user')
            password = dom_user.get('password')
            host_aliases = get_host_aliases(pm_hostname, pm_alias, pm_ip, user)
            output_one_block(host_aliases, pm_ip, user,
                             password=password, ip=pm_ip, arch=pm_arch)


def dump_ssh_cfg_block_vm(dom_pms, dom_vms, dom_users):
    vms_num = len(dom_vms)
    vms_hostname_width = 3 if len(str(vms_num)) < 3 else len(str(vms_num))

    for vm in dom_vms:
        pm_id = vm.get('pmid')
        pm = get_pm_byid(dom_pms, pm_id)
        pm_ip = pm.get('ip')
        pm_arch = pm.get('arch')

        vm_iid = vm.get('iid')
        vm_ip = vm.get('ip')
        vm_alias = vm.get('alias')
        vm_os = vm.get('os')
        vm_role = vm.get('role')
        vm_id = vm_iid - DOMAIN_VM_IID_BASE
        vm_hostname = f"VM{vm_id:0{vms_hostname_width}}"

        for dom_user in dom_users:
            user = dom_user.get('user')
            password = dom_user.get('password')
            if vm_os.find('Anolis OS') == -1 and user == 'anuser':
                continue
            host_aliases = get_host_aliases(vm_hostname, vm_alias, vm_ip, user)
            output_one_block(host_aliases, vm_ip, user,
                             password=password, os=vm_os, role=vm_role, ip=pm_ip, arch=pm_arch)


def main(argc, argv):
    dom_pms = load_yaml(DOMAIN_PMS_YAML)
    dom_vms = load_yaml(DOMAIN_VMS_YAML)
    dom_users = load_yaml(DOMAIN_USERS_YAML)
    dump_ssh_cfg_block_vm(dom_pms, dom_vms, dom_users)
    dump_ssh_cfg_block_pm(dom_pms, dom_users)

    return 0


if __name__ == '__main__':
    sys.exit(main(len(sys.argv), sys.argv))
